package com.flink.hbase;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.*;
import java.util.concurrent.TimeUnit;

/**
 * JDBC 维度数据源
 * 定期从数据库读取维度数据并发送到下游
 */
public class JdbcDimensionSource extends RichSourceFunction<DimensionData> {

    private static final Logger LOG = LoggerFactory.getLogger(JdbcDimensionSource.class);
    
    private final String jdbcUrl;
    private final String username;
    private final String password;
    private final String query;
    private final long intervalMs;
    
    private transient Connection connection;
    private transient PreparedStatement statement;
    private volatile boolean isRunning = true;

    public JdbcDimensionSource(String jdbcUrl, String username, String password, 
                              String query, long intervalMs) {
        this.jdbcUrl = jdbcUrl;
        this.username = username;
        this.password = password;
        this.query = query;
        this.intervalMs = intervalMs;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        
        LOG.info("正在初始化 JDBC 连接: {}", jdbcUrl);
        
        // 加载 JDBC 驱动
        Class.forName("com.mysql.cj.jdbc.Driver");
        
        // 创建数据库连接
        connection = DriverManager.getConnection(jdbcUrl, username, password);
        connection.setAutoCommit(false);
        
        // 预编译查询语句
        statement = connection.prepareStatement(query);
        
        LOG.info("JDBC 连接初始化完成");
    }

    @Override
    public void run(SourceContext<DimensionData> ctx) throws Exception {
        while (isRunning) {
            try {
                // 执行查询
                LOG.debug("执行维度数据查询: {}", query);
                ResultSet resultSet = statement.executeQuery();
                
                int count = 0;
                while (resultSet.next() && isRunning) {
                    DimensionData dimensionData = new DimensionData();
                    
                    // 从结果集中提取数据
                    dimensionData.setDimensionKey(resultSet.getString("dimension_key"));
                    dimensionData.setDimensionValue(resultSet.getString("dimension_value"));
                    dimensionData.setCategory(resultSet.getString("category"));
                    dimensionData.setUpdateTime(resultSet.getTimestamp("update_time").getTime());
                    dimensionData.setIsActive(resultSet.getBoolean("is_active"));
                    
                    // 发送数据到下游
                    synchronized (ctx.getCheckpointLock()) {
                        ctx.collect(dimensionData);
                    }
                    count++;
                }
                
                resultSet.close();
                LOG.info("成功读取 {} 条维度数据", count);
                
                // 等待下次执行
                if (isRunning) {
                    TimeUnit.MILLISECONDS.sleep(intervalMs);
                }
                
            } catch (SQLException e) {
                LOG.error("执行 JDBC 查询失败", e);
                
                // 尝试重新连接
                try {
                    reconnect();
                } catch (Exception reconnectException) {
                    LOG.error("重新连接失败", reconnectException);
                    throw reconnectException;
                }
                
                // 短暂等待后重试
                TimeUnit.SECONDS.sleep(5);
            }
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
        LOG.info("JDBC 数据源正在停止");
    }

    @Override
    public void close() throws Exception {
        super.close();
        
        if (statement != null) {
            statement.close();
        }
        
        if (connection != null) {
            connection.close();
        }
        
        LOG.info("JDBC 连接已关闭");
    }

    /**
     * 重新连接数据库
     */
    private void reconnect() throws SQLException {
        LOG.warn("正在尝试重新连接数据库");
        
        // 关闭现有连接
        try {
            if (statement != null) {
                statement.close();
            }
            if (connection != null) {
                connection.close();
            }
        } catch (SQLException e) {
            LOG.warn("关闭连接时发生异常", e);
        }
        
        // 重新建立连接
        connection = DriverManager.getConnection(jdbcUrl, username, password);
        connection.setAutoCommit(false);
        statement = connection.prepareStatement(query);
        
        LOG.info("数据库重新连接成功");
    }

    /**
     * 创建维度数据源的静态工厂方法
     */
    public static JdbcDimensionSource create(String jdbcUrl, String username, String password, 
                                           String tableName, long intervalMs) {
        String query = String.format(
            "SELECT dimension_key, dimension_value, category, update_time, is_active " +
            "FROM %s WHERE is_active = true ORDER BY update_time DESC", 
            tableName
        );
        
        return new JdbcDimensionSource(jdbcUrl, username, password, query, intervalMs);
    }

    /**
     * 创建自定义查询的维度数据源
     */
    public static JdbcDimensionSource createWithCustomQuery(String jdbcUrl, String username, 
                                                           String password, String customQuery, 
                                                           long intervalMs) {
        return new JdbcDimensionSource(jdbcUrl, username, password, customQuery, intervalMs);
    }
} 